{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Code to Evaluate: \n",
    "  - Proposed model\n",
    "  - single-task baselines\n",
    "  - fine-tuning baselines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note - Please set the appropriate dataset, model and checkpoint paths before using\n",
    "\n",
    "This notebook contains commands for CS->BDD->IDD and CS->IDD->BDD settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import os\n",
    "import importlib\n",
    "import time\n",
    "\n",
    "from PIL import Image\n",
    "from argparse import ArgumentParser\n",
    "\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.transforms import Compose, CenterCrop, Normalize, Resize\n",
    "from torchvision.transforms import ToTensor, ToPILImage\n",
    "\n",
    "from dataset import SYSU, CDD, PRCV # original label id gt annotations\n",
    "\n",
    "from models.erfnet_RA_parallel_CD import Net as Net_RAP # proposed model\n",
    "# from erfnet import ERFNet as ERFNet_ind # single-task models \n",
    "# from erfnet_ftp1 import Net as ERFNet_ft1 # 1st stage FT/FE (Eg. CS->BDD)\n",
    "# from erfnet_ftp2 import Net as ERFNet_ft2 # 2nd stage FT/FE (CS|BDD->IDD)\n",
    "\n",
    "from transform import Relabel, ToLabel, Colorize # modify IDD label ids if saving colour maps. otherwise its fine. \n",
    "from iouEval import iouEval, getColorEntry, Evaluator\n",
    "# from torchsummary import summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/a409_home/anaconda3/envs/zhangyu_pytorch/lib/python3.7/site-packages/torchvision/transforms/transforms.py:288: UserWarning: Argument interpolation should be of type InterpolationMode instead of int. Please, use InterpolationMode enum.\n",
      "  \"Argument interpolation should be of type InterpolationMode instead of int. \"\n"
     ]
    }
   ],
   "source": [
    "NUM_CHANNELS = 3\n",
    "NUMC_city = 20\n",
    "NUMC_bdd = 20\n",
    "NUMC_idd = 27\n",
    "\n",
    "image_transform = ToPILImage()\n",
    "input_transform = Compose([\n",
    "    Resize([512,1024], Image.BILINEAR),\n",
    "    ToTensor(),\n",
    "])\n",
    "target_transform_cityscapes = Compose([\n",
    "    Resize([512,1024],  Image.NEAREST),\n",
    "    ToLabel(),\n",
    "    Relabel(255, NUMC_city-1),   \n",
    "])\n",
    "target_transform_IDD = Compose([\n",
    "    Resize([512,1024],  Image.NEAREST),\n",
    "    ToLabel(),\n",
    "    Relabel(255, NUMC_idd-1),  \n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pass dataset name, get val_loader, criterion with suitable weight\n",
    "def criterion_fn(data_name='cityscapes'): \n",
    "    CS_datadir = '/ssd_scratch/cvit/prachigarg/cityscapes/'\n",
    "    BDD_datadir = '/ssd_scratch/cvit/prachigarg/bdd100k/seg/'\n",
    "    IDD_datadir = '/ssd_scratch/cvit/prachigarg/IDD_Segmentation/'\n",
    "    \n",
    "    SYSU_datadir = '/home/a409/users/wenglean/my_progs/data/SYSU-CD'\n",
    "    CDD_datadir = '/home/a409/users/wenglean/my_progs/data/CDD/ChangeDetectionDataset/Real/subset'\n",
    "    PRCV_datadir = '/home/a409/users/wenglean/my_progs/data/prcv2021'\n",
    "\n",
    "    co_transform = None\n",
    "    co_transform_val = None\n",
    "    if data_name == 'SYSU':\n",
    "        print('taking SYSU')\n",
    "        dataset_train = SYSU(SYSU_datadir, co_transform, 'train')\n",
    "#         dataset_val = SYSU(SYSU_datadir, co_transform_val, 'val')\n",
    "        dataset_val = SYSU(SYSU_datadir, co_transform_val, 'val')\n",
    "    elif data_name == 'CDD':\n",
    "        print('taking CDD')\n",
    "        dataset_train = CDD(CDD_datadir, co_transform, 'train')\n",
    "        dataset_val = CDD(CDD_datadir, co_transform_val, 'val')\n",
    "    elif data_name == 'PRCV':\n",
    "        print('taking PRCV')\n",
    "        dataset_train = PRCV(PRCV_datadir, co_transform, 'train')\n",
    "        dataset_val = PRCV(PRCV_datadir, co_transform_val, 'val')\n",
    "\n",
    "    loader_val = DataLoader(dataset_val, num_workers=4,\n",
    "                            batch_size=1, shuffle=False)\n",
    "\n",
    "#     weight = weight.cuda()\n",
    "#     criterion = nn.CrossEntropyLoss(weight=weight)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    \n",
    "    return loader_val, criterion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval(model, dataset_loader, criterion, task, num_classes):\n",
    "    model.eval()\n",
    "    epoch_loss_val = []\n",
    "    num_cls = num_classes[task]\n",
    "    evaluator_cd = Evaluator(num_class=2)  # class1:change class2:not change\n",
    "    evaluator_cd.reset()\n",
    "#     iouEvalVal = iouEval(num_cls, num_cls-1)\n",
    "\n",
    "    with torch.no_grad():\n",
    "#         for step, (images, labels, filename, filenameGt) in enumerate(dataset_loader):\n",
    "            # inputs size: torch.Size([1, 20, 512, 1024])\n",
    "        for step, sampled_batch in enumerate(dataset_loader):\n",
    "            image_A, image_B, labels = sampled_batch['image'][0], sampled_batch['image'][1], sampled_batch['label']\n",
    "            images = torch.cat((image_A, image_B), 1)\n",
    "                \n",
    "            start_time = time.time()\n",
    "            inputs = images.cuda()\n",
    "            targets = labels.long().cuda()\n",
    "\n",
    "            outputs = model(inputs, task)\n",
    "\n",
    "            loss = criterion(outputs, targets)\n",
    "            epoch_loss_val.append(loss.item())\n",
    "\n",
    "#             iouEvalVal.addBatch(outputs.max(1)[1].unsqueeze(1).data, targets.unsqueeze(1).data)\n",
    "            evaluator_cd.add_batch(targets, outputs.max(1)[1])\n",
    "#     average_epoch_loss_val = sum(epoch_loss_val) / len(epoch_loss_val)\n",
    "#     iouVal = 0\n",
    "#     iouVal, iou_classes = iouEvalVal.getIoU()\n",
    "#     p = iouEvalVal.getPre()\n",
    "#     r = iouEvalVal.getRecall()\n",
    "#     f1 = iouEvalVal.getF1()\n",
    "#     print('check val fn, loss, acc: ', iouVal) \n",
    "\n",
    "        mIoU_cd = evaluator_cd.Mean_Intersection_over_Union()\n",
    "        Precision_cd = evaluator_cd.Precision()\n",
    "        Recall_cd = evaluator_cd.Recall()\n",
    "        F1_cd = evaluator_cd.F1()\n",
    "        \n",
    "#     return iou_classes, iouVal, p, r, f1\n",
    "    return mIoU_cd, Precision_cd, Recall_cd, F1_cd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "taking SYSU\n",
      "total 12000 sample\n",
      "total 4000 sample\n",
      "taking CDD\n",
      "total 10000 sample\n",
      "total 2998 sample\n",
      "taking PRCV\n",
      "total 6000 sample\n",
      "total 2000 sample\n"
     ]
    }
   ],
   "source": [
    "loader_val_SYSU, criterion_SYSU = criterion_fn('SYSU')\n",
    "loader_val_CDD, criterion_CDD = criterion_fn('CDD')\n",
    "loader_val_PRCV, criterion_PRCV = criterion_fn('PRCV')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### KLD-RAPFT-dlr2 - PROPOSED RESULTS "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hi, inside erfnet_RA_parallel 0 1\n",
      "0.754664 tensor(0.7807) tensor(0.7758) tensor(0.7782)\n"
     ]
    }
   ],
   "source": [
    "# STEP 1 SYSU (with the DS-RAP and DS-BN)\n",
    "model_step1 = Net_RAP([2], 1, 0) \n",
    "model_step1 = torch.nn.DataParallel(model_step1).cuda()\n",
    "saved_model = torch.load('/home/a409/users/wenglean/my_progs/save/SYSU/model_best_SYSU_erfnet_RA_parallel_50_6RAP_FT_step1.pth.tar')\n",
    "model_step1.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "# iou_classes_step1_SYSU, val_acc_step1_SYSU, p, r, f1 = eval(model_step1,loader_val_SYSU, criterion_SYSU, 0, [2])\n",
    "val_acc_step1_SYSU, p, r, f1 = eval(model_step1,loader_val_SYSU, criterion_SYSU, 0, [2])\n",
    "print(val_acc_step1_SYSU, p, r, f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hi, inside erfnet_RA_parallel 1 2\n",
      "0.69653624 tensor(0.7996) tensor(0.6247) tensor(0.7014)\n",
      "0.60862494 tensor(0.7691) tensor(0.3397) tensor(0.4713)\n"
     ]
    }
   ],
   "source": [
    "# STEP 2 SYSU->CDD \n",
    "model_step2 = Net_RAP([2, 2], 2, 1) \n",
    "model_step2 = torch.nn.DataParallel(model_step2).cuda()\n",
    "saved_model = torch.load('/home/a409/users/wenglean/my_progs/save/SYSU2CDD/model_best_CDD_erfnet_RA_parallel_50_6ours-SYSU1-CDD2_step2.pth.tar')\n",
    "model_step2.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou, p, r, f1 = eval(model_step2, loader_val_SYSU, criterion_SYSU, 0, [2, 2])\n",
    "print(iou, p, r, f1)\n",
    "\n",
    "iou, p, r, f1 = eval(model_step2, loader_val_CDD, criterion_CDD, 1, [2, 2])\n",
    "print(iou, p, r, f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hi, inside erfnet_RA_parallel 1 2\n",
      "0.6959068 tensor(0.7446) tensor(0.6748) tensor(0.7080)\n",
      "0.69523275 tensor(0.7301) tensor(0.5678) tensor(0.6388)\n"
     ]
    }
   ],
   "source": [
    "# STEP 2 SYSU->CDD \n",
    "model_step2 = Net_RAP([2, 2], 2, 1) \n",
    "model_step2 = torch.nn.DataParallel(model_step2).cuda()\n",
    "saved_model = torch.load('/home/a409/users/wenglean/my_progs/save/SYSU2CDD/checkpoint_CDD_erfnet_RA_parallel_50_6ours-SYSU1-CDD2_step2.pth.tar')\n",
    "model_step2.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou, p, r, f1 = eval(model_step2, loader_val_SYSU, criterion_SYSU, 0, [2, 2])\n",
    "print(iou, p, r, f1)\n",
    "\n",
    "iou, p, r, f1 = eval(model_step2, loader_val_CDD, criterion_CDD, 1, [2, 2])\n",
    "print(iou, p, r, f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hi, inside erfnet_RA_parallel 1 2\n",
      "tensor(0.6458, dtype=torch.float64)\n",
      "tensor(0.5911, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# STEP 2 CS->IDD\n",
    "model_step2 = Net_RAP([20, 27], 2, 1) \n",
    "model_step2 = torch.nn.DataParallel(model_step2).cuda()\n",
    "saved_model = torch.load('/home2/prachigarg/temp_MDILSS/checkpoints/RAP_FT_KLD/step2idd/model_best_IDD_erfnet_RA_parallel_150_6RAP_FT_dlr2-5e-6-KLD-ouput-1e-1-IDD_step2.pth.tar')\n",
    "model_step2.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou_classes_step2_CS, val_acc_step2_CS = eval(model_step2, loader_val_CS, criterion_CS, 0, [20, 27])\n",
    "print(val_acc_step2_CS)\n",
    "\n",
    "iou_classes_step2_IDD, val_acc_step2_IDD = eval(model_step2, loader_val_IDD, criterion_IDD, 1, [20, 27])\n",
    "print(val_acc_step2_IDD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hi, inside erfnet_RA_parallel 2 3\n",
      "0.7036772 tensor(0.7581) tensor(0.6793) tensor(0.7165)\n",
      "0.6106379 tensor(0.6348) tensor(0.3942) tensor(0.4864)\n",
      "0.68106896 tensor(0.7966) tensor(0.5997) tensor(0.6843)\n"
     ]
    }
   ],
   "source": [
    "# STEP 3 SYSU|CDD->PRCV\n",
    "model_step3 = Net_RAP([2, 2, 2], 3, 2)\n",
    "model_step3 = torch.nn.DataParallel(model_step3).cuda()\n",
    "saved_model = torch.load('/home/a409/users/wenglean/my_progs/save/SYSU2CDD2PRCV/model_best_PRCV_erfnet_RA_parallel_50_6OURS-SYSU1-CDD2-PRCV3_step3.pth.tar')\n",
    "model_step3.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou, p, r, f1 = eval(model_step3, loader_val_SYSU, criterion_SYSU, 0, [2, 2, 2])\n",
    "print(iou, p, r, f1)\n",
    "\n",
    "iou, p, r, f1 = eval(model_step3, loader_val_CDD, criterion_CDD, 1, [2, 2, 2])\n",
    "print(iou, p, r, f1)\n",
    "\n",
    "iou, p, r, f1 = eval(model_step3, loader_val_PRCV, criterion_PRCV, 2, [2, 2, 2])\n",
    "print(iou, p, r, f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hi, inside erfnet_RA_parallel 2 3\n",
      "0.6959598 tensor(0.7472) tensor(0.6720) tensor(0.7076)\n",
      "0.61485004 tensor(0.5652) tensor(0.4569) tensor(0.5053)\n",
      "0.63088715 tensor(0.7580) tensor(0.5156) tensor(0.6137)\n"
     ]
    }
   ],
   "source": [
    "# STEP 3 SYSU|CDD->PRCV checkpoint\n",
    "model_step3 = Net_RAP([2, 2, 2], 3, 2)\n",
    "model_step3 = torch.nn.DataParallel(model_step3).cuda()\n",
    "saved_model = torch.load('/home/a409/users/wenglean/my_progs/save/SYSU2CDD2PRCV/checkpoint_PRCV_erfnet_RA_parallel_50_6OURS-SYSU1-CDD2-PRCV3_step3.pth.tar')\n",
    "model_step3.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou, p, r, f1 = eval(model_step3, loader_val_SYSU, criterion_SYSU, 0, [2, 2, 2])\n",
    "print(iou, p, r, f1)\n",
    "\n",
    "iou, p, r, f1 = eval(model_step3, loader_val_CDD, criterion_CDD, 1, [2, 2, 2])\n",
    "print(iou, p, r, f1)\n",
    "\n",
    "iou, p, r, f1 = eval(model_step3, loader_val_PRCV, criterion_PRCV, 2, [2, 2, 2])\n",
    "print(iou, p, r, f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hi, inside erfnet_RA_parallel 2 3\n",
      "tensor(0.6255, dtype=torch.float64)\n",
      "tensor(0.5385, dtype=torch.float64)\n",
      "tensor(0.5590, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# STEP 3 CS|IDD->BDD\n",
    "model_step3 = Net_RAP([20, 27, 20], 3, 2) \n",
    "model_step3 = torch.nn.DataParallel(model_step3).cuda()\n",
    "saved_model = torch.load('/home2/prachigarg/temp_MDILSS/checkpoints/RAP_FT_KLD/CS1_IDD2_BDD3/checkpoint_BDD_erfnet_RA_parallel_150_6OURS-CS1_IDD2_BDD3_step3.pth.tar')\n",
    "model_step3.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou_classes_step3_CS, val_acc_step3_CS = eval(model_step3, loader_val_CS, criterion_CS, 0, [20, 27, 20])\n",
    "print(val_acc_step3_CS)\n",
    "\n",
    "iou_classes_step3_IDD, val_acc_step3_IDD = eval(model_step3, loader_val_IDD, criterion_IDD, 1, [20, 27, 20])\n",
    "print(val_acc_step3_IDD)\n",
    "\n",
    "iou_classes_step3_BDD, val_acc_step3_BDD = eval(model_step3, loader_val_BDD, criterion_BDD, 2, [20, 27, 20])\n",
    "print(val_acc_step3_BDD)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Single-Task Baselines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.7255, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# STEP 1, single task\n",
    "# Cityscapes\n",
    "model_step1 = ERFNet_ind(20) \n",
    "model_step1 = torch.nn.DataParallel(model_step1).cuda()\n",
    "saved_model = torch.load('/home2/prachigarg/temp_MDILSS/checkpoints/single-task/model_best_cityscapes_prenc.pth.tar')\n",
    "model_step1.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou_classes_ST_cs, val_acc_ST_cs = eval(model_step1, loader_val_CS, criterion_CS, 0, [20])\n",
    "print(val_acc_ST_cs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.5410, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# BDD100k\n",
    "saved_model = torch.load('/home2/prachigarg/temp_MDILSS/checkpoints/single-task/checkpoint_BDD_prenc.pth.tar')\n",
    "model_step1.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou_classes_ST_bdd, val_acc_ST_bdd = eval(model_step1, loader_val_BDD, criterion_BDD, 0, [20])\n",
    "print(val_acc_ST_bdd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.6197, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# IDD\n",
    "model_step1 = ERFNet_ind(27)\n",
    "model_step1 = torch.nn.DataParallel(model_step1).cuda()\n",
    "saved_model = torch.load('/home2/prachigarg/temp_MDILSS/checkpoints/single-task/checkpoint_IDD_prenc.pth.tar')\n",
    "model_step1.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou_classes_ST_idd, val_acc_ST_idd = eval(model_step1, loader_val_IDD, criterion_IDD, 0, [27])\n",
    "print(val_acc_ST_idd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finetune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_finetune(model, dataset_loader, criterion, task, num_classes):\n",
    "    model.eval()\n",
    "    epoch_loss_val = []\n",
    "    num_cls = num_classes[task]\n",
    "\n",
    "    iouEvalVal = iouEval(num_cls, num_cls-1)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for step, (images, labels, filename, filenameGt) in enumerate(dataset_loader):\n",
    "            # inputs size: torch.Size([1, 20, 512, 1024])\n",
    "            start_time = time.time()\n",
    "            inputs = images.cuda()\n",
    "            targets = labels.cuda()\n",
    "            if task == 0:\n",
    "                outputs = model(inputs, decoder_old=True, decoder_new=False)\n",
    "            elif task == 1:\n",
    "                outputs = model(inputs, decoder_old=False, decoder_new=True)\n",
    "\n",
    "            loss = criterion(outputs, targets[:, 0])\n",
    "            epoch_loss_val.append(loss.item())\n",
    "\n",
    "            iouEvalVal.addBatch(outputs.max(1)[1].unsqueeze(1).data, targets.data)\n",
    "\n",
    "    average_epoch_loss_val = sum(epoch_loss_val) / len(epoch_loss_val)\n",
    "\n",
    "    iouVal = 0\n",
    "    iouVal, iou_classes = iouEvalVal.getIoU()\n",
    "\n",
    "    print('check val fn, loss, acc: ', iouVal) \n",
    "        \n",
    "    return iou_classes, iouVal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_finetune3(model, dataset_loader, criterion, task, num_classes=[20, 20, 27]):\n",
    "    model.eval()\n",
    "    epoch_loss_val = []\n",
    "    num_cls = num_classes[task]\n",
    "\n",
    "    iouEvalVal = iouEval(num_cls, num_cls-1)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for step, (images, labels, filename, filenameGt) in enumerate(dataset_loader):\n",
    "            # inputs size: torch.Size([1, 20, 512, 1024])\n",
    "            start_time = time.time()\n",
    "            inputs = images.cuda()\n",
    "            targets = labels.cuda()\n",
    "            if task == 0:\n",
    "                outputs = model(inputs, decoder_old1=True, decoder_old2=False, decoder_new=False)\n",
    "            elif task == 1:\n",
    "                outputs = model(inputs, decoder_old1=False, decoder_old2=True, decoder_new=False)\n",
    "            elif task == 2:\n",
    "                outputs = model(inputs, decoder_old1=False, decoder_old2=False, decoder_new=True)\n",
    "\n",
    "            loss = criterion(outputs, targets[:, 0])\n",
    "            epoch_loss_val.append(loss.item())\n",
    "\n",
    "            iouEvalVal.addBatch(outputs.max(1)[1].unsqueeze(1).data, targets.data)\n",
    "\n",
    "    average_epoch_loss_val = sum(epoch_loss_val) / len(epoch_loss_val)\n",
    "\n",
    "    iouVal = 0\n",
    "    iouVal, iou_classes = iouEvalVal.getIoU()\n",
    "\n",
    "    print('check val fn, loss, acc: ', iouVal) \n",
    "        \n",
    "    return iou_classes, iouVal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "check val fn, loss, acc:  tensor(0.4005, dtype=torch.float64)\n",
      "tensor(0.4005, dtype=torch.float64)\n",
      "check val fn, loss, acc:  tensor(0.5274, dtype=torch.float64)\n",
      "tensor(0.5274, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Step 2: Finetune CS model on BDD\n",
    "model_step2 = ERFNet_ft1(num_classes_old=20, num_classes_new=20)\n",
    "model_step2 = torch.nn.DataParallel(model_step2).cuda()\n",
    "saved_model = torch.load('/home2/prachigarg/temp_MDILSS/checkpoints/FineTune/checkpoint_erfnet_ftp1_150_6_Finetune-CStoBDD-final.pth.tar')\n",
    "model_step2.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou_classes_step2_FT1_CS, val_acc_step2_FT1_CS = eval_finetune(model_step2, loader_val_CS, criterion_CS, 0, [20, 20])\n",
    "print(val_acc_step2_FT1_CS)\n",
    "\n",
    "iou_classes_step2_FT1_BDD, val_acc_step2_FT1_BDD = eval_finetune(model_step2, loader_val_BDD, criterion_BDD, 1, [20, 20])\n",
    "print(val_acc_step2_FT1_BDD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "check val fn, loss, acc:  tensor(0.3049, dtype=torch.float64)\n",
      "tensor(0.3049, dtype=torch.float64)\n",
      "check val fn, loss, acc:  tensor(0.3205, dtype=torch.float64)\n",
      "tensor(0.3205, dtype=torch.float64)\n",
      "check val fn, loss, acc:  tensor(0.6103, dtype=torch.float64)\n",
      "tensor(0.6103, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Step 3: Finetune CS|BDD model on IDD\n",
    "model_step3 = ERFNet_ft2(20, 20, 27) \n",
    "model_step3 = torch.nn.DataParallel(model_step3).cuda()\n",
    "saved_model = torch.load('/home2/prachigarg/temp_MDILSS/checkpoints/FineTune/model_best_erfnet_ftp2_150_6_Finetune-code-CSBDDtoIDD-FT.pth.tar')\n",
    "model_step3.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou_classes_step3_CS, val_acc_step3_CS = eval_finetune3(model_step3, loader_val_CS, criterion_CS, 0, [20, 20, 27])\n",
    "print(val_acc_step3_CS)\n",
    "\n",
    "iou_classes_step3_BDD, val_acc_step3_BDD = eval_finetune3(model_step3, loader_val_BDD, criterion_BDD, 1, [20, 20, 27])\n",
    "print(val_acc_step3_BDD)\n",
    "\n",
    "iou_classes_step3_IDD, val_acc_step3_IDD = eval_finetune3(model_step3, loader_val_IDD, criterion_IDD, 2, [20, 20, 27])\n",
    "print(val_acc_step3_IDD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "check val fn, loss, acc:  tensor(0.3619, dtype=torch.float64)\n",
      "tensor(0.3619, dtype=torch.float64)\n",
      "check val fn, loss, acc:  tensor(0.2630, dtype=torch.float64)\n",
      "tensor(0.2630, dtype=torch.float64)\n",
      "check val fn, loss, acc:  tensor(0.5337, dtype=torch.float64)\n",
      "tensor(0.5337, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "# Step 3: Finetune CS|IDD model on BDD\n",
    "model_step3 = ERFNet_ft2(20, 27, 20) \n",
    "model_step3 = torch.nn.DataParallel(model_step3).cuda()\n",
    "saved_model = torch.load('/home2/prachigarg/temp_MDILSS/checkpoints/FineTune/model_best_erfnet_ftp2_150_6_FT_CS1_IDD2_BDD3.pth.tar')\n",
    "model_step3.load_state_dict(saved_model['state_dict'])\n",
    "\n",
    "iou_classes_step3_CS, val_acc_step3_CS = eval_finetune3(model_step3, loader_val_CS, criterion_CS, 0, [20, 27, 20])\n",
    "print(val_acc_step3_CS)\n",
    "\n",
    "iou_classes_step3_IDD, val_acc_step3_IDD = eval_finetune3(model_step3, loader_val_IDD, criterion_IDD, 1, [20, 27, 20])\n",
    "print(val_acc_step3_IDD)\n",
    "\n",
    "iou_classes_step3_BDD, val_acc_step3_BDD = eval_finetune3(model_step3, loader_val_BDD, criterion_BDD, 2, [20, 27, 20])\n",
    "print(val_acc_step3_BDD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
